module PdbParser

using Printf
using Utils
using Structure

export gen_pdb_from_file, to_pdb
export PDB, PdbModel

struct PdbModel
    modelid::Int32
    items::Vector{GeneralAtom}
end

struct PDB
    title::String
    crystallographic::Vector{Float64}
    models::Vector{PdbModel}
    connectivity::Vector{Vector{Int32}}
end

function gen_pdb_from_file(filename::AbstractString)
    println("Reading pdb file $filename...")
    lines = readlines(filename)
    # 处理title
    title = findfirst(x -> startswith(x, "REMARK"), lines) !== nothing ? strip(lines[findfirst(x -> startswith(x, "REMARK"), lines)][7:end]) : ""
    # 处理盒子
    crystal_lines = filter(startswith.("CRYST1"), lines)
    crystal = Vector{Int32}()
    if length(crystal_lines) > 0
        crystal = parse.(Float64, split(crystal_lines[1])[2:4])
    end
    # 处理坐标
    model_index = findall(x -> match(r"^(MODEL\s*\d+|ENDMDL)$", x) !== nothing, lines)
    models = Vector{PdbModel}()
    if length(model_index) <= 2     # 只有一个model
        push!(models, PdbModel(1, collect_atoms(lines)))
    else                            # 有多个model
        for i in 1:length(model_index) / 2
            push!(models, PdbModel(i, collect_atoms(lines)))
        end
    end
    # TODO: 处理CONECT
    connect_lines = filter(startswith.("CONECT"), lines)
    connect = []
    if length(connect_lines) > 0
        for line in connect_lines
            push!(connect, parse.(Int32, split(line[7:end])))
        end
    end
    finish_hint("Finished reading $filename.\n")
    return PDB(title, crystal, models, connect)
end

function collect_atoms(lines::Vector{String})
    items = Vector{GeneralAtom}()
    for line in lines
        if startswith(line, "ATOM") || startswith(line, "HETATM")
            push!(items, gen_pdb_atom_from_line(line))
        elseif startswith(line, "TER")
            push!(items, gen_pdb_ter_from_line(line))
        end
    end
    return items
end

function gen_pdb_atom_from_line(atomline::AbstractString)
    typ = strip(atomline[1:6])
    atnum = parse(Int32, strip(atomline[7:11]))
    atid = atid_for_human(strip(atomline[13:16]))
    altloc = atomline[17:17]
    resname = strip(atomline[18:20])
    chainid = atomline[22:22] == " " ? "X" : atomline[22:22]
    resid = parse(Int32, strip(atomline[23:26]))
    icode = atomline[27:27]
    x = parse(Float64, strip(atomline[31:38])) / 10     # Å to nm
    y = parse(Float64, strip(atomline[39:46])) / 10     # Å to nm
    z = parse(Float64, strip(atomline[47:54])) / 10     # Å to nm
    occupy = parse(Float64, strip(atomline[55:60]))
    bf = parse(Float64, strip(atomline[61:66]))
    atname = ""
    charge = ""
    try
        atname = strip(atomline[77:78])
    catch
    end
    if isempty(atname)
        atname = string(atid[1])
        warning_hint("Warning: missing element, will be guessed as $atname from atom id ($atid)\n")
    end
    try
        charge = strip(atomline[79:80])
    catch
    end
    return GeneralAtom(; typ=typ, atid=atnum, atname=atid, altloc=altloc, resname=resname, chainname=chainid, resid=resid, 
        icode=icode, x=x, y=y, z=z, occupy=occupy, bf=bf, element=atname, charge_formal=charge)
end

function gen_pdb_ter_from_line(terline::AbstractString)
    typ = terline[1:3]
    atnum = 0
    resname = ""
    chainid = ""
    resid = 0
    icode = ""
    try
        atnum = parse(Int32, strip(terline[7:11]))
        resname = strip(terline[18:20])
        chainid = terline[22:22] == " " ? "X" : terline[22:22]
        resid = parse(Int32, strip(terline[23:26]))
        icode = terline[27:27]
    catch
    end
    return GeneralAtom(; typ=typ, atid=atnum, resname=resname, chainname=chainid, resid=resid, icode=icode)
end

function atid_for_human(atid::AbstractString)
    # 从pdb使用的atid到方便理解的atid
    if length(atid) == 4
        # 理论上应该是HG11, HG12这样, 核酸可能是H2'1, H2'2
        if match(r"[1-9][A-Z][ABGDEZH0-9][1-9']", atid) !== nothing
            return atid[2:end] * atid[1]
        else
            warning_hint("Warning: atom name may be unstandard: $atid\n")
            return atid
        end
    else
        return strip(atid)   # strip 因为 atid_for_pdb 会加一个空格
    end
end

function to_pdb(pdb::PDB, pdb_file::AbstractString)
    pdb_file = open(pdb_file, "w")
    write(pdb_file, "REMARK   Created by SuperMDA (https://gitee.com/supernova_bingbing/SuperMDA)\nREMARK   " * pdb.title * "\n")
    # TODO: 写box信息
    for model in pdb.models
        write(pdb_file, "MODEL      " * "$(model.modelid)\n")
        for item in model.items
            write(pdb_file, to_str_pdb(item))
        end
        write(pdb_file, "ENDMDL\n")
    end
    for connect in pdb.connectivity
        write(pdb_file, "CONECT")
        for i in connect
            write(pdb_file, @sprintf("%5d", i))
        end
        write(pdb_file, "\n")
    end
    write(pdb_file, "END\n")
    close(pdb_file)
end

end # module
